%% 幻觉控制刺激优化版 - 只生成PNG图像
clear; sca

%% 设置参数
areaSum = 1200; 
number_sets = 1:15; 
image_iter = 20; 
color_num = 50; 
img_size = [224, 224]; % 图像尺寸

% 预先生成颜色
colors = rand(color_num, 3) * 0.8 + 0.2;

% 创建主输出文件夹
main_output_dir = 'illusion_control1_images';
if ~exist(main_output_dir, 'dir')
    mkdir(main_output_dir);
end

% 预计算圆的半径，避免重复计算
circle_radii = zeros(length(number_sets), 1);
for tt = 1:length(number_sets)
    nn = number_sets(tt);
    syms circle_radius positive;
    circle_radius = solve((pi*circle_radius^2)*nn*2 + nn*2*circle_radius == areaSum, circle_radius);
    circle_radii(tt) = double(circle_radius);
end

% 检查是否可以使用并行计算
can_use_parallel = exist('parpool', 'file') > 0;
if can_use_parallel
    try
        p = gcp('nocreate');
        if isempty(p)
            parpool('local');
        end
    catch
        can_use_parallel = false;
    end
end

% 进度显示
total_images = color_num * length(number_sets) * image_iter;
fprintf('开始生成 %d 张图像...\n', total_images);
progress_counter = 0;
progress_interval = max(1, floor(total_images/20)); % 每完成5%显示一次

% 预生成圆形掩码模板
circle_masks = cell(length(circle_radii), 1);
for tt = 1:length(circle_radii)
    r = circle_radii(tt);
    [X, Y] = meshgrid(-ceil(r):ceil(r), -ceil(r):ceil(r));
    mask = (X.^2 + Y.^2 <= r^2);
    circle_masks{tt} = mask;
end

% 主循环 - 使用并行处理加速（如果可用）
if can_use_parallel
    parfor cc = 1:color_num
        process_color(cc, colors, circle_radii, circle_masks, number_sets, image_iter, main_output_dir, img_size);
    end
else
    for cc = 1:color_num
        process_color(cc, colors, circle_radii, circle_masks, number_sets, image_iter, main_output_dir, img_size);
        
        % 更新进度
        progress_counter = progress_counter + length(number_sets) * image_iter;
        if mod(progress_counter, progress_interval) == 0
            fprintf('进度: %.1f%% 完成\n', 100 * progress_counter / total_images);
        end
    end
end

fprintf('所有图像生成完成!\n');

%% 辅助函数
function process_color(cc, colors, circle_radii, circle_masks, number_sets, image_iter, main_output_dir, img_size)
    current_color = colors(cc, :);
    
    % 计算并调整背景颜色
    bg_color = calculate_background_color(current_color);
    
    % 创建颜色文件夹
    color_diff = round(sqrt(sum((round(current_color * 255) - round(bg_color * 255)).^2)));
    color_folder_name = sprintf('color_%d_%.3f_%.3f_%.3f_diff%.0f', cc, current_color(1), current_color(2), current_color(3), color_diff);
    color_folder_path = fullfile(main_output_dir, color_folder_name);
    if ~exist(color_folder_path, 'dir')
        mkdir(color_folder_path);
    end
    
    % 处理每种数量的圆
    for tt = 1:length(number_sets)
        nn = number_sets(tt);
        circle_radius = circle_radii(tt);
        circle_mask = circle_masks{tt};
        
        for ii = 1:image_iter
            % 创建背景图像
            img = ones(img_size) .* reshape(bg_color, [1, 1, 3]);
            
            % 生成圆和连接线
            [img, success] = generate_circles_and_lines(img, nn, circle_radius, circle_mask, current_color, bg_color, img_size);
            
            if success
                % 保存图像
                img_filename = sprintf('illusion_nn%d_img%d.png', nn, ii);
                img_path = fullfile(color_folder_path, img_filename);
                imwrite(img, img_path);
            else
                % 如果生成失败，重试
                ii = ii - 1;
            end
        end
    end
end

function bg_color = calculate_background_color(current_color)
    % 使用HSV色彩空间计算背景颜色
    hsv_color = rgb2hsv(current_color);
    hsv_color(2) = hsv_color(2) * 0.4;   % 降低饱和度到40%
    hsv_color(3) = min(1, hsv_color(3) * 1.8);  % 增加明度，但不超过1
    bg_color = hsv2rgb(hsv_color);
    
    % 检查颜色差值，确保差值>100 (0-255范围内)
    stim_color_255 = round(current_color * 255);
    bg_color_255 = round(bg_color * 255);
    color_diff = sqrt(sum((stim_color_255 - bg_color_255).^2));
    
    % 调整背景色直到满足差值要求
    adjustment_attempts = 0;
    while color_diff < 100 && adjustment_attempts < 10
        if hsv_color(3) < 0.9
            hsv_color(3) = min(1, hsv_color(3) + 0.1);
        else
            hsv_color(2) = max(0, hsv_color(2) - 0.1);
        end
        bg_color = hsv2rgb(hsv_color);
        bg_color_255 = round(bg_color * 255);
        color_diff = sqrt(sum((stim_color_255 - bg_color_255).^2));
        adjustment_attempts = adjustment_attempts + 1;
    end
    
    % 如果差值仍不足，使用默认高对比度背景
    if color_diff < 100
        luminance = 0.2126 * current_color(1) + 0.7152 * current_color(2) + 0.0722 * current_color(3);
        if luminance > 0.5
            bg_color = [0.1, 0.1, 0.1]; % 暗色背景
        else
            bg_color = [0.9, 0.9, 0.9]; % 亮色背景
        end
    end
end

function [img, success] = generate_circles_and_lines(img, nn, circle_radius, circle_mask, circle_color, line_color, img_size)
    % 初始化
    all_centers_x = [];
    all_centers_y = [];
    connections = [];
    locmtmp = [];
    success = true;
    max_attempts = 50; % 最大尝试次数
    
    % 生成第一对连接圆
    padding = ceil(5*circle_radius);
    width = img_size(2);
    height = img_size(1);
    
    % 第一个圆
    x1 = padding + randi(width - 2*padding);
    y1 = padding + randi(height - 2*padding);
    
    % 计算第二个圆的位置
    theta = 2*pi*rand;
    distance = 3*circle_radius;
    x2 = round(x1 + distance * cos(theta));
    y2 = round(y1 + distance * sin(theta));
    
    % 检查是否在边界内
    if x2 < padding || x2 > width-padding || y2 < padding || y2 > height-padding
        % 尝试另一个方向
        attempts = 0;
        while (x2 < padding || x2 > width-padding || y2 < padding || y2 > height-padding) && attempts < max_attempts
            theta = 2*pi*rand;
            x2 = round(x1 + distance * cos(theta));
            y2 = round(y1 + distance * sin(theta));
            attempts = attempts + 1;
        end
        
        if attempts >= max_attempts
            success = false;
            return;
        end
    end
    
    % 计算中点位置
    locm = complex(0.5*(x1+x2), 0.5*(y1+y2));
    
    % 保存参数
    all_centers_x = [x1, x2];
    all_centers_y = [y1, y2];
    connections = [1, 2];
    locmtmp = locm;
    
    % 生成剩余的连接圆
    attempts = 0;
    pair_attempts = 0;
    
    while length(locmtmp) < nn && pair_attempts < max_attempts*2
        % 重置当前对尝试计数
        attempts = 0;
        pair_ok = false;
        
        while attempts < max_attempts && ~pair_ok
            % 第一个圆
            x1 = padding + randi(width - 2*padding);
            y1 = padding + randi(height - 2*padding);
            
            % 第二个圆
            theta = 2*pi*rand;
            x2 = round(x1 + distance * cos(theta));
            y2 = round(y1 + distance * sin(theta));
            
            % 检查边界
            if x2 >= padding && x2 <= width-padding && y2 >= padding && y2 <= height-padding
                % 计算中点
                locm = complex(0.5*(x1+x2), 0.5*(y1+y2));
                
                % 检查与现有中点的距离
                distancestmp = abs(locm - locmtmp);
                min_dist = 6*circle_radius;
                if all(distancestmp > min_dist)
                    pair_ok = true;
                end
            end
            
            attempts = attempts + 1;
        end
        
        if pair_ok
            idx1 = length(all_centers_x) + 1;
            idx2 = length(all_centers_x) + 2;
            all_centers_x = [all_centers_x, x1, x2];
            all_centers_y = [all_centers_y, y1, y2];
            connections = [connections; idx1, idx2];
            locmtmp = [locmtmp, locm];
        end
        
        pair_attempts = pair_attempts + 1;
    end
    
    if length(locmtmp) < nn
        success = false;
        return;
    end
    
    % 绘制所有圆点
    for i = 1:length(all_centers_x)
        img = draw_circle_fast(img, round(all_centers_x(i)), round(all_centers_y(i)), circle_radius, circle_mask, circle_color);
    end
    
    % 绘制连接线
    for c = 1:size(connections, 1)
        idx1 = connections(c, 1);
        idx2 = connections(c, 2);
        x1 = round(all_centers_x(idx1));
        y1 = round(all_centers_y(idx1));
        x2 = round(all_centers_x(idx2));
        y2 = round(all_centers_y(idx2));
        img = draw_line_fast(img, x1, y1, x2, y2, line_color);
    end
end

function img = draw_circle_fast(img, cx, cy, r, mask, color)
    % 使用预计算的掩码快速绘制圆
    r_int = ceil(r);
    mask_size = 2*r_int+1;
    
    % 计算图像中应用掩码的范围
    y_start = max(1, cy-r_int);
    y_end = min(size(img,1), cy+r_int);
    x_start = max(1, cx-r_int);
    x_end = min(size(img,2), cx+r_int);
    
    % 计算掩码中应用的范围
    mask_y_start = 1 + max(0, 1-(cy-r_int));
    mask_y_end = mask_size - max(0, (cy+r_int)-size(img,1));
    mask_x_start = 1 + max(0, 1-(cx-r_int));
    mask_x_end = mask_size - max(0, (cx+r_int)-size(img,2));
    
    % 截取有效区域
    mask_sub = mask(mask_y_start:mask_y_end, mask_x_start:mask_x_end);
    
    % 应用颜色到每个通道
    for c = 1:3
        temp = img(y_start:y_end, x_start:x_end, c);
        temp(mask_sub) = color(c);
        img(y_start:y_end, x_start:x_end, c) = temp;
    end
end

function img = draw_line_fast(img, x1, y1, x2, y2, color)
    % 使用bresenham算法绘制直线
    dx = abs(x2 - x1);
    dy = abs(y2 - y1);
    steep = dy > dx;
    
    if steep
        [x1, y1] = deal(y1, x1);
        [x2, y2] = deal(y2, x2);
    end
    
    if x1 > x2
        [x1, x2] = deal(x2, x1);
        [y1, y2] = deal(y2, y1);
    end
    
    dx = x2 - x1;
    dy = abs(y2 - y1);
    err = dx / 2;
    
    if y1 < y2
        ystep = 1;
    else
        ystep = -1;
    end
    
    for x = x1:x2
        if steep
            if x >= 1 && x <= size(img, 1) && y1 >= 1 && y1 <= size(img, 2)
                img(x, y1, :) = reshape(color, [1, 1, 3]);
            end
        else
            if x >= 1 && x <= size(img, 2) && y1 >= 1 && y1 <= size(img, 1)
                img(y1, x, :) = reshape(color, [1, 1, 3]);
            end
        end
        
        err = err - dy;
        if err < 0
            y1 = y1 + ystep;
            err = err + dx;
        end
    end
end